/*
elliptic curve asymmetric
*/

package sm234

import (
	"crypto/elliptic"
	"errors"
	"io"
	"math/big"
	"sync"
)

// PublicKey represents an SM2 public key.
type PublicKey struct {
	curve elliptic.Curve
	X, Y  *big.Int
}

// PrivateKey represents a SM2 private key.
type PrivateKey struct {
	PublicKey
	D *big.Int
}

var initonce sync.Once
var p256_sm2 *elliptic.CurveParams

func get_p256_sm2() elliptic.Curve {
	initonce.Do(func() {
		// See FIPS 186-3, section D.2.4
		p256_sm2 = &elliptic.CurveParams{Name: "SM2-P-256"}
		p256_sm2.P, _ = new(big.Int).SetString("FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF", 16)
		p256_sm2.N, _ = new(big.Int).SetString("FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123", 16)
		p256_sm2.B, _ = new(big.Int).SetString("28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93", 16)
		p256_sm2.Gx, _ = new(big.Int).SetString("32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7", 16)
		p256_sm2.Gy, _ = new(big.Int).SetString("BC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0", 16)
		p256_sm2.BitSize = 256
	})
	return p256_sm2
}

func randFieldElement(c elliptic.Curve, rand io.Reader) (*big.Int, error) {
	params := c.Params()
	one := big.NewInt(1)

	b := make([]byte, params.BitSize/8+8)
	_, err := io.ReadFull(rand, b)
	if err != nil {
		return nil, err
	}

	val := new(big.Int).SetBytes(b)
	n := new(big.Int).Sub(params.N, one)
	n = n.Sub(n, one) //n-2

	// 1 <= k <= n-2
	val.Mod(val, n)
	val.Add(val, one)
	return val, nil
}

func SM2GenerateKey(rand io.Reader) (*PrivateKey, error) {
	c := get_p256_sm2()
	k, err := randFieldElement(c, rand)
	if err != nil {
		return nil, err
	}

	priv := new(PrivateKey)
	priv.D = k
	priv.PublicKey.curve = c
	priv.PublicKey.X, priv.PublicKey.Y = c.ScalarBaseMult(k.Bytes())
	return priv, nil
}

func SM2BytesToKey(vbs []byte) (*PrivateKey, error) {
	c := get_p256_sm2()
	k := new(big.Int).SetBytes(vbs)

	priv := new(PrivateKey)
	priv.D = k
	priv.PublicKey.curve = c
	priv.PublicKey.X, priv.PublicKey.Y = c.ScalarBaseMult(k.Bytes())

	return priv, nil
}

// digest : 32 bytes
func (priv *PrivateKey) Sign(rand io.Reader, digest []byte) (r, s *big.Int, err error) {
	var k, e *big.Int

	//
	c := priv.PublicKey.curve
	N := c.Params().N
	if N.Sign() == 0 {
		return nil, nil, errors.New("invalid elliptic curve params(N)")
	}

	// msg hash to bigInt
	one := big.NewInt(1)
	d1inv := new(big.Int).Add(priv.D, one)
	d1inv.ModInverse(d1inv, N)
	e = new(big.Int).SetBytes(digest)

	for {

		for {

			k, err = randFieldElement(c, rand)
			if err != nil {
				r, s = nil, nil
				return
			}

			r, _ = c.ScalarBaseMult(k.Bytes())
			r.Add(r, e)
			r.Mod(r, N)
			if r.Sign() == 0 {
				continue
			}

			t := new(big.Int).Add(r, k)
			if t.Cmp(N) != 0 {
				break
			}
		}

		//
		s = new(big.Int).Mul(r, priv.D)
		s.Sub(k, s)
		s.Mod(s, N)
		s.Mul(s, d1inv)
		s.Mod(s, N)
		if s.Sign() != 0 {
			break
		}
	}

	// finish...
	return
}

// digest : 32 bytes
func (publ *PublicKey) Verify(digest []byte, r, s *big.Int) bool {
	c := publ.curve
	N := c.Params().N

	if (r.Sign() <= 0) || (s.Sign() <= 0) {
		return false
	}

	if (r.Cmp(N) >= 0) || (s.Cmp(N) >= 0) {
		return false
	}

	//
	t := new(big.Int).Add(r, s)
	t.Mod(t, N)
	if t.Sign() == 0 {
		return false
	}

	//
	x1, y1 := c.ScalarBaseMult(s.Bytes())
	x2, y2 := c.ScalarMult(publ.X, publ.Y, t.Bytes())
	x, _ := c.Add(x1, y1, x2, y2)

	//
	e := new(big.Int).SetBytes(digest)
	e.Add(e, x)
	e.Mod(e, N)
	return e.Cmp(r) == 0
}
